{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\nCompile ONNX Models\n===================\n**Author**: `Joshua Z. Zhang <https://zhreshold.github.io/>`_\n\nThis article is an introductory tutorial to deploy ONNX models with Relay.\n\nFor us to begin with, ONNX package must be installed.\n\nA quick solution is to install protobuf compiler, and\n\n.. code-block:: bash\n\n    pip install --user onnx onnxoptimizer\n\nor please refer to official site.\nhttps://github.com/onnx/onnx\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import onnx\nimport numpy as np\nimport tvm\nfrom tvm import te\nimport tvm.relay as relay\nfrom tvm.contrib.download import download_testdata"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Load pretrained ONNX model\n---------------------------------------------\nThe example super resolution model used here is exactly the same model in onnx tutorial\nhttp://pytorch.org/tutorials/advanced/super_resolution_with_caffe2.html\nwe skip the pytorch model construction part, and download the saved onnx model\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "model_url = \"\".join(\n    [\n        \"https://gist.github.com/zhreshold/\",\n        \"bcda4716699ac97ea44f791c24310193/raw/\",\n        \"93672b029103648953c4e5ad3ac3aadf346a4cdc/\",\n        \"super_resolution_0.2.onnx\",\n    ]\n)\nmodel_path = download_testdata(model_url, \"super_resolution.onnx\", module=\"onnx\")\n# now you have super_resolution.onnx on disk\nonnx_model = onnx.load(model_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Load a test image\n---------------------------------------------\nA single cat dominates the examples! This model takes a single input image of size\n224x224 and outputs a scaled image that is 3x greater than the input along each\naxis, a 672x672 image. Re-scale the cat image to fit this input shape then\nconvert to `YCbCr`. The super resolution model will then be applied to the\nluminance (`Y`) channel.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from PIL import Image\n\nimg_url = \"https://github.com/dmlc/mxnet.js/blob/main/data/cat.png?raw=true\"\nimg_path = download_testdata(img_url, \"cat.png\", module=\"data\")\nimg = Image.open(img_path).resize((224, 224))\nimg_ycbcr = img.convert(\"YCbCr\")  # convert to YCbCr\nimg_y, img_cb, img_cr = img_ycbcr.split()\nx = np.array(img_y)[np.newaxis, np.newaxis, :, :]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Compile the model with relay\n---------------------------------------------\nTypically ONNX models mix model input values with parameter values, with\nthe input having the name `1`. This model dependent, and you should check\nwith the documentation for your model to determine the full input and\nparameter name space.\n\nPassing in the shape dictionary to the `relay.frontend.from_onnx` method\ntells relay which ONNX parameters are inputs, and which are parameters, and\nprovides a static definition of the input size.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "target = \"llvm\"\n\ninput_name = \"1\"\nshape_dict = {input_name: x.shape}\nmod, params = relay.frontend.from_onnx(onnx_model, shape_dict)\n\nwith tvm.transform.PassContext(opt_level=1):\n    executor = relay.build_module.create_executor(\n        \"graph\", mod, tvm.cpu(0), target, params\n    ).evaluate()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Execute on TVM\n---------------------------------------------\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "dtype = \"float32\"\ntvm_output = executor(tvm.nd.array(x.astype(dtype))).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Display results\n---------------------------------------------\nWe put input and output image neck to neck. The luminance channel, `Y` is the output\nfrom the model. The chroma channels `Cb` and `Cr` are resized to match with a simple\nbicubic algorithm. The image is then recombined and converted back to `RGB`.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from matplotlib import pyplot as plt\n\nout_y = Image.fromarray(np.uint8((tvm_output[0, 0]).clip(0, 255)), mode=\"L\")\nout_cb = img_cb.resize(out_y.size, Image.BICUBIC)\nout_cr = img_cr.resize(out_y.size, Image.BICUBIC)\nresult = Image.merge(\"YCbCr\", [out_y, out_cb, out_cr]).convert(\"RGB\")\ncanvas = np.full((672, 672 * 2, 3), 255)\ncanvas[0:224, 0:224, :] = np.asarray(img)\ncanvas[:, 672:, :] = np.asarray(result)\nplt.imshow(canvas.astype(np.uint8))\nplt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Notes\n---------------------------------------------\nBy default, ONNX defines models in terms of dynamic shapes. The ONNX importer\nretains that dynamism upon import, and the compiler attempts to convert the model\ninto a static shapes at compile time. If this fails, there may still be dynamic\noperations in the model. Not all TVM kernels currently support dynamic shapes,\nplease file an issue on discuss.tvm.apache.org if you hit an error with dynamic kernels.\n\nThis particular model was build using an older version of ONNX. During the import\nphase ONNX importer will run the ONNX verifier, which may throw a `Mismatched attribute type`\nwarning. Because TVM supports a number of different ONNX versions, the Relay model\nwill still be valid.\n\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}